package com.bw.test1;

import org.apache.flink.types.Row;

import com.alibaba.alink.operator.batch.BatchOperator;
import com.alibaba.alink.operator.batch.regression.LinearRegPredictBatchOp;
import com.alibaba.alink.operator.batch.regression.LinearRegTrainBatchOp;
import com.alibaba.alink.operator.batch.source.MemSourceBatchOp;
import org.junit.Test;

import java.util.Arrays;
import java.util.List;

public class LinearRegTrainBatchOpTest {
    @Test
    public void testLinearRegTrainBatchOp() throws Exception {

        BatchOperator.setParallelism(1);
        // 1.准备数据源
        List<Row> df = Arrays.asList(
                Row.of(2, 1, 1),
                Row.of(3, 2, 1),
                Row.of(4, 3, 2),
                Row.of(2, 4, 1),
                Row.of(2, 2, 1),
                Row.of(4, 3, 2),
                Row.of(1, 2, 1),
                Row.of(5, 3, 3)
        );

        // 2、选择特征和label
        BatchOperator<?> batchData = new MemSourceBatchOp(df, "f0 int, f1 int, label int");

        // 训练组件
        BatchOperator<?> lr = new LinearRegTrainBatchOp()
                .setFeatureCols("f0", "f1")
                .setLabelCol("label")
                .lazyPrintModelInfo();

        // 3. 训练模型
        BatchOperator model = batchData.link(lr);

        // 4、预测
        // 预测组件
        BatchOperator<?> predictor = new LinearRegPredictBatchOp()
                .setPredictionCol("pred");
        predictor.linkFrom(model, batchData).print();
    }
}